#include "data_dumper.h"

#include <sys/stat.h>
#include <sys/types.h>
#include <unistd.h>

#include <fstream>
#include <iostream>
#include <iterator>
#include <map>
#include <numeric>
#include <sstream>
#include <vector>
#include <thread>

#include "file_utils.h"

#define UNLIKELY(x) __builtin_expect(!!(x), 0)
#define LIKELY(x) __builtin_expect(!!(x), 1)

std::vector<uint8_t> MindSporeOpRangeData::encode() {
  std::vector<uint8_t> tlvBytes;
  size_t totalSize = tlvBufferLength + op_name.size();
  tlvBytes.reserve(totalSize);

  // Type
  uint16_t dataType = static_cast<uint16_t>(tag);
  encodeFixedData<uint16_t>(dataType, tlvBytes);

  // Length
  encodeFixedData<uint32_t>(totalSize, tlvBytes);

  // Fixed data
  // 5 * 8 = 40 bytes
  encodeFixedData<uint64_t>(thread_id, tlvBytes);
  encodeFixedData<uint64_t>(flow_id, tlvBytes);
  encodeFixedData<uint64_t>(step, tlvBytes);
  encodeFixedData<uint64_t>(start_time_ns, tlvBytes);
  encodeFixedData<uint64_t>(end_time_ns, tlvBytes);
  // 4 bytes
  encodeFixedData<int32_t>(process_id, tlvBytes);
  // 2 * 2 = 4 bytes
  encodeFixedData<uint16_t>(module_index, tlvBytes);
  encodeFixedData<uint16_t>(event_index, tlvBytes);
  encodeFixedData<uint16_t>(stage_index, tlvBytes);
  // 1 byte
  encodeFixedData<int8_t>(level, tlvBytes);
  // 1 * 3 = 3 bytes
  encodeFixedData<bool>(is_async, tlvBytes);
  encodeFixedData<bool>(is_graph_data, tlvBytes);
  encodeFixedData<bool>(is_stage, tlvBytes);

  encodeStrData(static_cast<uint16_t>(OpRangeDataType::NAME), op_name, tlvBytes);
  return tlvBytes;
}


// std::vector<uint8_t> MindSporeOpRangeData::encode() {
//   std::vector<uint8_t> tmp;
//   std::vector<uint8_t> tlvBytes;

//   // Fixed data
//   // 5 * 8 = 40 bytes
//   encodeFixedData<uint64_t>(thread_id, tmp);
//   encodeFixedData<uint64_t>(flow_id, tmp);
//   encodeFixedData<uint64_t>(step, tmp);
//   encodeFixedData<uint64_t>(start_time_ns, tmp);
//   encodeFixedData<uint64_t>(end_time_ns, tmp);
//   // 4 bytes
//   encodeFixedData<int32_t>(process_id, tmp);
//   // 2 * 2 = 4 bytes
//   encodeFixedData<uint16_t>(module_index, tmp);
//   encodeFixedData<uint16_t>(event_index, tmp);
//   encodeFixedData<uint16_t>(stage_index, tmp);
//   // 1 byte
//   encodeFixedData<int8_t>(level, tmp);
//   // 1 * 3 = 3 bytes
//   encodeFixedData<bool>(is_async, tmp);
//   encodeFixedData<bool>(is_graph_data, tmp);
//   encodeFixedData<bool>(is_stage, tmp);

//   encodeStrData(static_cast<uint16_t>(OpRangeDataType::NAME), op_name, tmp);

//   uint16_t dataType = static_cast<uint16_t>(OpRangeDataType::NAME);
//   for (size_t i = 0; i < sizeof(uint16_t); ++i) {
//     tlvBytes.push_back((dataType >> (i * 8)) & 0xff);
//   }
//   uint32_t length = tmp.size();
//   for (size_t i = 0; i < sizeof(uint32_t); ++i) {
//     tlvBytes.push_back((length >> (i * 8)) & 0xff);
//   }
//   tlvBytes.insert(tlvBytes.end(), tmp.cbegin(), tmp.cend());
//   return tlvBytes;
// }

DataDumper::DataDumper() : path_(""), start_(false), init_(false) {}

DataDumper &DataDumper::GetInstance() {
  static DataDumper instance;
  return instance;
}
DataDumper::~DataDumper() { UnInit(); }

void DataDumper::Init(const std::string &path, size_t capacity) {
  path_ = path;
  data_chunk_buf_.Init(capacity);
  init_.store(true);
}

void DataDumper::UnInit() {
  if (init_.load()) {
    data_chunk_buf_.UnInit();
    init_.store(false);
    start_.store(false);
    for (auto &f : fd_map_) {
      if (f.second != nullptr) {
        fclose(f.second);
        f.second = nullptr;
      }
    }
    fd_map_.clear();
  }
}

void DataDumper::Start() {
  if (!init_.load() || Thread::Start() != 0) {
    return;
  }
  start_.store(true);
}

void DataDumper::Stop() {
  if (start_.load() == true) {
    start_.store(false);
    Thread::Stop();
  }
  Flush();
  std::cout << "Dump finished, total size: " << dump_count_  << " bytes" << std::endl;
}

void DataDumper::GatherAndDumpData() {
  std::unordered_map<std::string, std::vector<uint8_t>> dataMap;
  uint64_t batchSize = 0;
  while (batchSize < kBatchMaxLen) {
    std::unique_ptr<BaseReportData> data{nullptr};
    if (UNLIKELY(!data_chunk_buf_.Pop(data) || data == nullptr)) {
      break;
    }
    std::vector<uint8_t> encodeData = data->encode();
    ++batchSize;
    const std::string &key =
        kReportFileTypeMap.at(static_cast<ReportFileType>(data->tag));
    auto iter = dataMap.find(key);
    if (iter == dataMap.end()) {
      dataMap.insert({key, encodeData});
    } else {
      iter->second.insert(iter->second.end(), encodeData.cbegin(),
                          encodeData.cend());
    }
  }
  if (dataMap.size() > 0) {
    static bool create_flag = true;
    if (create_flag) {
      create_flag = !Utils::CreateDir(this->path_);
    }
    Dump(dataMap);
  }
}

void DataDumper::Run() {
  for (;;) {
    if (!start_.load()) {
      break;
    }
    if (data_chunk_buf_.Size() > kNotifyInterval) {
      GatherAndDumpData();
    } else {
      usleep(kMaxWaitTimeUs);
      //std::this_thread::sleep_for(std::chrono::milliseconds(kMaxWaitTimeUs));
    }
  }
}

void DataDumper::Flush() {
  while (data_chunk_buf_.Size() != 0) {
    GatherAndDumpData();
  }
}

void DataDumper::Report(std::unique_ptr<BaseReportData> data) {
  if (UNLIKELY(!start_.load() || data == nullptr)) {
    return;
  }
  data_chunk_buf_.Push(std::move(data));
}

void DataDumper::Dump(
    const std::unordered_map<std::string, std::vector<uint8_t>> &dataMap) {
  for (auto &data : dataMap) {
    FILE *fd = nullptr;
    const std::string dump_file = path_ + "/" + data.first;
    auto iter = fd_map_.find(dump_file);
    if (iter == fd_map_.end()) {
      if (!Utils::IsFileExist(dump_file) && !Utils::CreateFile(dump_file)) {
        continue;
      }
      fd = fopen(dump_file.c_str(), "ab");
      if (fd == nullptr) {
        continue;
      }
      fd_map_.insert({dump_file, fd});
    } else {
      fd = iter->second;
    }
    // std::cout << "dump file: " << dump_file << ", size: " << data.second.size()
    //           << std::endl;
    fwrite(reinterpret_cast<const char *>(data.second.data()), sizeof(char),
           data.second.size(), fd);
    dump_count_ += data.second.size();
  }
}